#!/usr/bin/env python3
import os
import hashlib
import subprocess
import tempfile
import uuid
import json
import time
import threading
from pathlib import Path
from queue import Queue
from typing import List, Dict
from datetime import datetime, timedelta

from flask import (
    Flask, render_template, request, jsonify,
    send_from_directory, abort, send_file,
)
from faster_whisper import WhisperModel
from pydub import AudioSegment
from pydub.generators import Sine

# ────── 路径常量 ──────
BASE = Path(__file__).resolve().parent
UPLOADS = BASE / "uploads"
OUTPUTS = BASE / "outputs"
SUBS = BASE / "subtitles"
KEYWORDS_DIR = BASE / "keywords"
for d in (UPLOADS, OUTPUTS, SUBS):
    d.mkdir(exist_ok=True)

# ────── 配置 ──────
MODEL_NAME   = "turbo"
CPU_THREADS  = 8
SAMPLE_RATE  = 16000
SEG_LEN_SEC  = 30
PRE_PAD_MS, POST_PAD_MS = 50, 100
BEEP_HZ = 1000
CLEANUP_THRESHOLD = 3 * 24 * 60 * 60  # 3 天

def load_keywords(folder: Path):
    kws, cmap = set(), {}
    for p in folder.glob("*.*"):
        if "-" not in p.stem:
            continue
        bad, _ = p.stem.split("-", 1)
        bad = bad.strip()
        kws.add(bad)
        cmap[bad] = str(p.resolve())
    return kws, cmap

KEYWORDS, CUSTOM_MAP = load_keywords(KEYWORDS_DIR)
app = Flask(__name__)

def cleanup_stale_files():
    now = time.time()
    for folder in (UPLOADS, OUTPUTS):
        for f in folder.iterdir():
            if not f.is_file() or f.name == "EMPTY":
                continue
            if now - f.stat().st_mtime > CLEANUP_THRESHOLD:
                try:
                    f.unlink()
                except:
                    pass

def md5_file(path: Path, buf: int = 1 << 20) -> str:
    h = hashlib.md5()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(buf), b""):
            h.update(chunk)
    return h.hexdigest()

def transcribe_or_load(wav: Path, cache_json: Path) -> List[Dict]:
    if cache_json.exists():
        return json.loads(cache_json.read_text("utf-8"))
    model = WhisperModel(
        MODEL_NAME, device="cpu",
        cpu_threads=CPU_THREADS, compute_type="int8"
    )
    seg_dir = wav.parent / "seg"
    seg_dir.mkdir(exist_ok=True)
    subprocess.run([
        "ffmpeg", "-y", "-i", str(wav),
        "-f", "segment", "-segment_time", str(SEG_LEN_SEC),
        "-reset_timestamps", "1",
        str(seg_dir / "seg_%04d.wav")
    ], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    words, offset = [], 0.0
    for seg in sorted(seg_dir.glob("seg_*.wav")):
        dur = float(subprocess.check_output([
            "ffprobe", "-v", "error",
            "-show_entries", "format=duration",
            "-of", "default=nokey=1:noprint_wrappers=1",
            str(seg)
        ]))
        segs, _ = model.transcribe(
            str(seg), language="zh", word_timestamps=True
        )
        for sg in segs:
            for w in sg.words:
                txt = w.word.strip()
                if txt:
                    words.append({
                        "word": txt,
                        "start": w.start + offset,
                        "end": w.end + offset
                    })
        offset += dur
    cache_json.write_text(json.dumps(words, ensure_ascii=False), "utf-8")
    return words

def match_segments(words: List[Dict]) -> List[tuple]:
    """
    针对内置 KEYWORDS 做最长匹配，返回 (start, end, phrase) 列表。
    """
    chars  = [w["word"] for w in words]
    starts = [w["start"] for w in words]
    ends   = [w["end"]   for w in words]
    maxL   = max((len(k) for k in KEYWORDS), default=1)
    res, i = [], 0
    while i < len(chars):
        hit = None
        for L in range(maxL, 0, -1):
            if i + L > len(chars): continue
            phrase = "".join(chars[i:i+L])
            if phrase in KEYWORDS:
                hit = (starts[i], ends[i+L-1], phrase)
                break
        if hit:
            res.append(hit)
            i += len(hit[2])
        else:
            i += 1
    return res

def replace_audio(orig_wav: Path, targets: List[tuple]) -> AudioSegment:
    orig = AudioSegment.from_wav(str(orig_wav))
    if not targets:
        return orig
    out, last = AudioSegment.empty(), 0
    for st, ed, kw, mode in sorted(targets, key=lambda x: x[0]):
        a = max(0, int(st*1000) - PRE_PAD_MS)
        b = int(ed*1000) + POST_PAD_MS
        out += orig[last:a]
        dur = b - a
        if mode == "silence":
            rep = AudioSegment.silent(duration=dur)
        elif mode == "beep":
            rep = Sine(BEEP_HZ).to_audio_segment(duration=dur)
        else:  # custom
            path = CUSTOM_MAP.get(kw)
            if path and Path(path).exists():
                clip = AudioSegment.from_file(path)
                if len(clip) < dur:
                    clip *= dur // len(clip) + 1
                rep = clip[:dur]
            else:
                rep = AudioSegment.silent(duration=dur)
        out += rep
        last = b
    out += orig[last:]
    return out

def pipeline(job: Dict):
    in_video     = job["in_video"]
    out_video    = job["out_video"]
    sub_json     = job["sub_json"]
    replace_list = job["replace"]
    silence_list = job["silence"]

    # 1. 提取音轨
    tmp_dir = tempfile.TemporaryDirectory()
    wav = Path(tmp_dir.name) / "audio.wav"
    subprocess.run([
        "ffmpeg", "-y", "-i", str(in_video),
        "-vn", "-ac", "1", "-ar", str(SAMPLE_RATE), str(wav)
    ], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

    # 2. 转写或加载缓存
    words = transcribe_or_load(wav, sub_json)

    # 3. 内置关键词匹配，筛出替换任务
    base_hits = match_segments(words)
    rep_hits = [(st, ed, kw) for st, ed, kw in base_hits if kw in replace_list]

    # 4. 全量消音匹配（包括用户新加词）
    sil_hits: List[tuple] = []
    for phrase in silence_list:
        L = len(phrase)
        for i in range(len(words) - L + 1):
            seq = "".join(w["word"] for w in words[i:i+L])
            if seq == phrase:
                st = words[i]["start"]
                ed = words[i+L-1]["end"]
                sil_hits.append((st, ed, phrase))

    # 5. 合并 targets 并标记模式
    targets: List[tuple] = []
    for st, ed, kw in rep_hits:
        targets.append((st, ed, kw, "custom"))
    for st, ed, kw in sil_hits:
        targets.append((st, ed, kw, "silence"))

    # 6. 生成新音轨并合成视频
    new_audio = replace_audio(wav, targets)
    new_wav = Path(tmp_dir.name) / "new.wav"
    new_audio.export(str(new_wav), format="wav")
    subprocess.run([
        "ffmpeg", "-y",
        "-i", str(in_video),
        "-i", str(new_wav),
        "-c:v", "copy", "-map", "0:v:0", "-map", "1:a:0",
        "-shortest", str(out_video)
    ], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    tmp_dir.cleanup()

    # 7. 更新作业状态
    jobs[job["id"]].update({
        "status": "done",
        "url": f"/download/{job['md5']}/{out_video.name}",
        "filename": out_video.name
    })

# ────── 队列机制 ──────
job_queue = Queue()
jobs: Dict[str, Dict] = {}

def worker():
    while True:
        job = job_queue.get()
        jid = job["id"]
        jobs[jid]["status"] = "processing"
        for idx, queued in enumerate(list(job_queue.queue)):
            jobs[queued["id"]]["position"] = idx
        try:
            pipeline(job)
        except Exception as e:
            jobs[jid].update({"status": "failed", "error": str(e)})
        job_queue.task_done()

threading.Thread(target=worker, daemon=True).start()

# ────── Flask 路由 ──────
@app.route("/")
def index():
    cleanup_stale_files()
    items = [{"display": p.stem, "file": p.name}
             for p in KEYWORDS_DIR.glob("*.*") if "-" in p.stem]
    items.sort(key=lambda x: x["display"])

    three_days_ago = time.time() - CLEANUP_THRESHOLD
    vids = []
    for f in OUTPUTS.iterdir():
        if not f.is_file(): continue
        if f.stat().st_mtime < three_days_ago: continue
        vids.append({
            "name": f.name,
            "created_at": datetime.fromtimestamp(f.stat().st_mtime).strftime("%Y-%m-%d %H:%M"),
            "url": f"/download/{f.stem}/{f.name}"
        })
    vids.sort(key=lambda x: x["created_at"], reverse=True)

    return render_template(
        "index.html",
        kw_items=items,
        processed_videos=vids
    )

@app.route('/favicon.ico')
def favicon():
    return send_file(os.path.join(app.root_path, 'favicon.ico'))

@app.route("/process", methods=["POST"])
def enqueue():
    if "file" not in request.files or not request.files["file"].filename:
        abort(400, "no file")
    f = request.files["file"]
    replace_list = json.loads(request.form.get("replace", "[]"))
    silence_list = json.loads(request.form.get("silence", "[]"))

    tmp = UPLOADS / f"tmp_{uuid.uuid4().hex}"
    f.save(str(tmp))
    md5 = md5_file(tmp)
    ext = Path(f.filename).suffix or ".mp4"
    upath = UPLOADS / f"{md5}{ext}"
    if not upath.exists():
        tmp.rename(upath)
    else:
        tmp.unlink(missing_ok=True)

    sub  = SUBS / f"{md5}.json"
    fname = f"（AI脱敏）{Path(f.filename).name}"
    outp  = OUTPUTS / fname
    if outp.exists():
        outp = OUTPUTS / f"（AI脱敏）{uuid.uuid4().hex[:6]}_{Path(f.filename).name}"

    job_id = uuid.uuid4().hex
    jobs[job_id] = {"status": "queued", "position": job_queue.qsize()}
    job = {
        "id": job_id,
        "md5": md5,
        "in_video": upath,
        "out_video": outp,
        "sub_json": sub,
        "replace": replace_list,
        "silence": silence_list
    }
    job_queue.put(job)
    return jsonify(job_id=job_id, position=jobs[job_id]["position"])

@app.route("/status/<job_id>")
def status(job_id):
    info = jobs.get(job_id)
    if not info:
        abort(404)
    return jsonify(info)

@app.route("/download/<md5>/<name>")
def download(md5, name):
    out_fp = OUTPUTS / name
    if not out_fp.exists():
        abort(404)
    return send_from_directory(str(OUTPUTS), name, as_attachment=True)

@app.route("/kw/<filename>")
def serve_kw_file(filename):
    file_path = KEYWORDS_DIR / filename
    if file_path.exists():
        return send_from_directory(str(KEYWORDS_DIR), filename)
    else:
        abort(404)

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5005, debug=False)
